#include <iostream>
#include <fstream>
#include <string>
#include <vector>
#include <map>

#include "CL/opencl.h"
#include "CL/cl.hpp"

static const size_t npos=std::string::npos;


void usage(char const* argv0,std::ostream& os=std::cout)
{
  os << "OpenCL offline compiler for AMD APP\n";
  os << "usage:" << argv0 << " [options] [*.cl]\n";
  os << "options:\n";
  os << "  --help               show this message\n";
  os << "  -o OUTPUTPREFIX      specify output filename prefix [default:cl_]\n";        
  os << "  --list-devices       list target device names and exit\n";
  os << "  --device=DEVICES     set build target devices\n";
  os << "                       DEVICES: comma separated list of DEVICE\n";
  os << "                       DEVICE: a GPU device name or a string \"CPU\"\n";
  os << "  --cl-opts=OPTSTRING  set option string for cl::Program::Build\n";
  os << std::endl;
}

bool loadSource(std::string* dest,std::string const& filename)
{
  std::ifstream ifs;
  ifs.open(filename.c_str(),ifs.binary);
  if (!ifs) return false;
  dest->clear();
  std::vector<char> buf(8192);
  while (!!ifs) {
    ifs.read(&buf[0],buf.size());
    dest->append(&buf[0],ifs.gcount());
  }
  return true;
}


std::string replaceSpaces(std::string const& s)
{
  std::string ret(s);
  for (size_t i=0;i!=s.length();++i) {
    if (ret[i]==' ') ret[i]='_';
  }
  return ret;
}

inline std::string getDeviceModifiedName(cl::Device const& dev)
{
  return replaceSpaces(dev.getInfo<CL_DEVICE_NAME>());
}

std::string getDeviceTargetName(cl::Device const& dev)
{
  if (dev.getInfo<CL_DEVICE_TYPE>()==CL_DEVICE_TYPE_CPU) {
    return "CPU";
  }
  return getDeviceModifiedName(dev);
}


bool saveBinaries(cl::Program const& prog,std::string const& prefix)
{
  std::vector<cl_device_id> devices(prog.getInfo<CL_PROGRAM_DEVICES>());
  std::vector<size_t> binary_sizes(prog.getInfo<CL_PROGRAM_BINARY_SIZES>());
  std::vector<char*> binaries(prog.getInfo<CL_PROGRAM_BINARIES>());
  for (size_t i=0;i!=binaries.size();++i) { 
    std::ofstream ofs;
    std::string filename;
    cl::Device device(devices[i]);
    filename=prefix+getDeviceTargetName(device)+".bin";
    ofs.open(filename.c_str(),ofs.binary);
    size_t size=binary_sizes[i];
    char* p=binaries[i];
    if (size) {ofs.write(p,size);}
    delete [] p;
  }
  return true;
}

bool getAMDPlatform(cl::Platform* dest)
{
  std::vector<cl::Platform> platforms;
  cl::Platform::get(&platforms);
  for (size_t i=0;i!=platforms.size();++i) {
    if (platforms[i].getInfo<CL_PLATFORM_VENDOR>()
	.find("Advanced Micro Devices")!=npos) {
      *dest=platforms[i];
      return true;
    }
  }
  return false;
}

void listDevices(std::vector<cl::Device> const& devices,std::ostream& os=std::cout)
{
  for (size_t i=0;i!=devices.size();++i) {
    cl::Device const& device(devices[i]);
    std::string name=getDeviceModifiedName(device);
    if (device.getInfo<CL_DEVICE_TYPE>()==CL_DEVICE_TYPE_CPU) {
      os << "CPU [" << name << "]\n";
    } else {
      os << name << "\n";
    }
  }
  os.flush();
}

bool getTargetDevices(std::vector<cl::Device>& devices,
		      std::vector<std::string> const& namelist)
{
  std::map<std::string,size_t> map;
  for (size_t i=0;i!=devices.size();++i) {
    cl::Device const& device(devices[i]);
    std::string name(getDeviceTargetName(device));
    map[name]=i;
  }

  bool ret=true;

  std::vector<cl::Device> temp;
  for (size_t i=0;i!=namelist.size();++i) {
    std::string const& name(namelist[i]);
    if (map.count(name)) {
      temp.push_back(devices[map[name]]);
    } else {
      std::cerr << "Device " << name << " is not found!!" << std::endl;
      ret=false;
    }
  }
  devices.swap(temp);
  return ret;
}


typedef std::vector<cl_context_properties> ContextProperties;

template <typename T>
ContextProperties& operator<<(ContextProperties& props,T const& value)
{
  props.push_back((cl_context_properties const&)value);
  return props;
}

int main(int argc,char* argv[])
{
  bool list=false;
  std::vector<std::string> devlist;
  std::string cloption;
  std::string ofile="cl_";

  char** p;
  for (p=&argv[1];*p;++p) {
    std::string s(*p);
    if (s=="--help" || s=="-h") {
      usage(argv[0]);
      exit(0);
    } else if (s=="--list-devices") {
      list=true;
    } else if (s.compare(0,9,"--device=")==0) {      
      for (size_t i=8,j=s.find_first_of(',',i+1);
	   i!=npos;
	   i=j,j=s.find_first_of(',',i+1)) {
	devlist.push_back(s.substr(i+1,j==npos?npos:j-i-1));
      }
    } else if (s.compare(0,10,"--cl-opts=")==0) {
      cloption+=s.substr(10)+" ";
    } else if (s == "-o") {
      ++p;
      if (*p) {
	ofile=*p;
      } else {
	std::cerr << "Option -o must be followed by filename prefix!!" << std::endl;
	exit(1);
      }
    } else if (s == "--") {
      ++p;
      break;
    } else {
      break;
    }
  }

  cl::Platform platform;
  if (!getAMDPlatform(&platform)) {
    std::cerr << "AMD platform is not found!!" << std::endl;
    exit(1);
  }

  ContextProperties cprops;

  cprops << CL_CONTEXT_PLATFORM << platform;
  cprops << CL_CONTEXT_OFFLINE_DEVICES_AMD << 1;
  cprops << 0;

  cl::Context ctxt(CL_DEVICE_TYPE_ALL,&cprops[0]);
  std::vector<cl::Device> devices(ctxt.getInfo<CL_CONTEXT_DEVICES>());
  if (list) {
    listDevices(devices);
    exit(0);
  }

  if (!getTargetDevices(devices,devlist)) {
    exit(1);
  }

  ctxt=cl::Context(devices,&cprops[0]);

  if (*p==0) {
    std::cerr << "No source files!!" << std::endl;
    exit(1);
  }

  std::vector<std::string> src_strs;
  while (*p) {
    src_strs.push_back(std::string());
    if (!loadSource(&src_strs.back(),*p)) {
      std::cerr << "File \""<< *p << "\" is not found!!" << std::endl;
      exit(1);
    }
    ++p;
  }

  cl::Program::Sources srcs;
  srcs.reserve(src_strs.size());
  for (size_t i=0;i!=src_strs.size();++i) {
    srcs.push_back(std::make_pair(src_strs[i].c_str(),src_strs[i].size()));
  }
  
  cl::Program prog(ctxt,srcs);
  cl_int err=prog.build(devices,cloption.c_str());
  for (size_t i=0;i!=devices.size();++i) {
    cl::Device const& device(devices[i]);
    std::string log(prog.getBuildInfo<CL_PROGRAM_BUILD_LOG>(device));
    if (log.empty()) continue;
    std::cerr << "[Device log " 
	      << replaceSpaces(device.getInfo<CL_DEVICE_NAME>()) 
	      << "]"
	      << std::endl;
    std::cerr << log << std::endl;
  }
  if (err != CL_SUCCESS) {
    exit(1);
  }
  
  saveBinaries(prog,ofile);
  exit(0);
}
